{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Open Vocabulary Object Segmentation using TAO Mask Grounding DINO\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## What is Mask Grounding DINO?\n",
    "\n",
    "Built on top of [Grounding DINO](https://arxiv.org/abs/2303.05499), Mask Grounding DINO is a NVIDIA proprietary open-set object segmentation model. Mask Grounding DINO can detect and segment arbitrary objects with human inputs such as category names or referring expressions. Compared to Grounding DINO, Mask Grounding DINO has an additional mask branch and mask prediction head.\n",
    "\n",
    "In TAO, only single type of backbone network is supported: [Swin](https://arxiv.org/abs/2103.14030). In this notebook, we use the pretrained Swin-Tiny Grounding DINO and showcase how we can finetune on COCO dataset for the state of the art mAP result.\n",
    "\n",
    "### Sample prediction of Swin-Tiny + Mask Grounding DINO model\n",
    "<img align=\"center\" src=\"sample.jpg\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained model and finetune a Mask Grounding DINO model on COCO dataset\n",
    "* Evaluate the trained model\n",
    "* Deploy the trained model\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Mask Grounding DINO using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate a trained model](#head-5)\n",
    "6. [Visualize inferences](#head-6)\n",
    "7. [Deploy](#head-7)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/grounding_dino/results`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"mask_grounding_dino\", \"results\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/mask_grounding_dino\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_DATA_DIR\n",
    "!mkdir -p $HOST_SPECS_DIR\n",
    "!mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": \"~/.cache\",\n",
    "           \"destination\": \"/.cache\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"64G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.10. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the COCO dataset for the tutorial. The following script will download COCO dataset automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "# Download the data\n",
    "!bash $HOST_SPECS_DIR/download_coco.sh $HOST_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verification\n",
    "!ls -l $HOST_DATA_DIR/raw-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create ODVG folder\n",
    "!mkdir -p $HOST_DATA_DIR/odvg\n",
    "!mkdir -p $HOST_DATA_DIR/odvg/annotations\n",
    "\n",
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert COCO to ODVG format required for Mask Grounding DINO\n",
    "!tao dataset annotations convert \\\n",
    "            -e $SPECS_DIR/convert.yaml \\\n",
    "            coco.ann_file=$DATA_DIR/raw-data/annotations/instances_train2017.json \\\n",
    "            results_dir=$DATA_DIR/odvg/annotations/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert COCO validation annotations to have categoy id ranging from 0 to 79.\n",
    "# This is required for computing validation loss during Mask Grounding DINO training.\n",
    "!tao dataset annotations convert \\\n",
    "            -e $SPECS_DIR/convert.yaml \\\n",
    "            coco.ann_file=$DATA_DIR/raw-data/annotations/instances_val2017.json \\\n",
    "            results_dir=$DATA_DIR/odvg/annotations/ \\\n",
    "            data.output_format=\"COCO\" \\\n",
    "            coco.use_all_categories=True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Download pre-trained model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/grounding_dino:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/grounding_dino:grounding_dino_swin_tiny_commercial_trainable_v1.0 --dest $LOCAL_PROJECT_DIR/mask_grounding_dino/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_PROJECT_DIR/mask_grounding_dino/grounding_dino_vgrounding_dino_swin_tiny_commercial_trainable_v1.0/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * train_data_sources:\n",
    "        * image_dir: the root directory for train images\n",
    "        * json_file: ODVG annotation file\n",
    "        * label_map: category id and category mapping\n",
    "    * val_data_sources: \n",
    "        * image_dir: the root directory for validation images\n",
    "        * json_file: annotation file for validation data. Required to be in COCO json format and the categoy id should be in the range of 0 ~ # of classes - 1\n",
    "    * max_labels: max number of positive + negative labels seen in a single batch. Larger max_labels usually result in better accuracy with longer training time.\n",
    "    * batch_size: batch size for dataloader\n",
    "    * workers: number of workers to do data loading\n",
    "* model: configure the model setting\n",
    "    * pretrained_backbone_path: path to the pretrained backbone model. Only Swin-variants are supported\n",
    "    * num_feature_levels: number of feature levels used from backbone\n",
    "    * dec_layers: number of decoder layers\n",
    "    * enc_layers: number of encoder layers\n",
    "    * num_queries: number of queries for the model\n",
    "    * num_select: number of top-k proposals to select from\n",
    "    * use_dn: flag to enable denoising during training\n",
    "    * dropout_ratio: drop out ratio\n",
    "* train: configure the training hyperparameters\n",
    "    * num_gpus: number of gpus \n",
    "    * num_nodes: number of nodes (num_nodes=1 for single node)\n",
    "    * val_interval: validation interval\n",
    "    * optim:\n",
    "        * lr_backbone: learning rate for backbone\n",
    "        * lr: learning rate for the rest of the model\n",
    "        * lr_steps: learning rate decay step milestone (MultiStep)\n",
    "    * num_epochs: number of epochs\n",
    "    * activation_checkpoint: recompute activations in the backward to save GPU memory. Default is `True`.\n",
    "    * precision: If set to fp16/bf16, the training is run on Automatic Mixed Precision (AMP)\n",
    "    * distributed_strategy: Default is `ddp`. `ddp_sharded` is also supported.\n",
    "\n",
    "Please refer to the TAO documentation about Grounding DINO to get all the parameters that are configurable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* Evaluation uses COCO metrics. For more info, please refer to: https://cocodataset.org/#detection-eval\n",
    "* *WARNING*: [according to the orirginal paper](https://arxiv.org/abs/2303.05499), COCO training was conducted using 8 A100 gpus. As a result, **we highly recommend that you run training with multiple high-end gpus (e.g. V100, A100)**\n",
    "* *Note*: The current training config should fit on a GPU with 16G of memory. Try to lower `batch_size` if you face OOM issue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change num_gpus in train.yaml based on your machine or update the NUM_TRAIN_GPUS env variable in the line below to the desired number of GPUs.\")\n",
    "os.environ[\"NUM_TRAIN_GPUS\"] = \"1\"\n",
    "\n",
    "!tao model mask_grounding_dino train \\\n",
    "           -e $SPECS_DIR/train.yaml \\\n",
    "           train.num_gpus=$NUM_TRAIN_GPUS \\\n",
    "           results_dir=$RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Trained checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/mask_gdino_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/mask_grounding_dino_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/mask_grounding_dino_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate a trained model <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this section, we run the `evaluate` tool to evaluate the trained model and produce the mAP metric.\n",
    "\n",
    "We provide evaluate.yaml specification files to configure the evaluate parameters including:\n",
    "\n",
    "* model: configure the model setting\n",
    "    * this config should remain same as your trained model's configuration.\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * test_data_sources:\n",
    "        * image_dir: the root directory for evaluatation images    \n",
    "        * json_file: Required to be in COCO json format and the categoy id should be in the range of 0 ~ # of classes - 1\n",
    "    * batch_size\n",
    "    * workers\n",
    "* evaluate:\n",
    "    * num_gpus: number of gpus\n",
    "    * conf_threshold: a threshold for confidence scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate on TAO model\n",
    "!tao model mask_grounding_dino evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/mask_grounding_dino_model.pth \\\n",
    "            results_dir=$RESULTS_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Visualize Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models and visualize the results. The `inference` tool produces annotated image outputs and txt files that contain prediction information.\n",
    "\n",
    "We provide inference.yaml specification files to configure the inference parameters including:\n",
    "\n",
    "* model: configure the model setting\n",
    "    * this config should remain same as your trained model's configuration\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * infer_data_sources:\n",
    "        * image_dir: the list of directories for inference images\n",
    "        * captions: list of phrases to run inference on. E.g. [\"person\", \"black cat\"]\n",
    "    * batch_size\n",
    "    * workers\n",
    "* inference\n",
    "    * conf_threshold: the confidence score threshold\n",
    "    * color_map: the color mapping for each phrase. The predicted bbox will be drawn with mapped color for each phrase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model mask_grounding_dino inference \\\n",
    "        -e $SPECS_DIR/infer.yaml \\\n",
    "        inference.checkpoint=$RESULTS_DIR/train/mask_grounding_dino_model.pth \\\n",
    "        results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"inference\", \"images_annotated\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy (**experimental**) <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "**Note:** The batch size of the exported ONNX model and TensorRT must be 1. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export the model to ONNX model\n",
    "!tao model mask_grounding_dino export \\\n",
    "           -e $SPECS_DIR/export.yaml \\\n",
    "           export.checkpoint=$RESULTS_DIR/train/mask_grounding_dino_model.pth \\\n",
    "           export.onnx_file=$RESULTS_DIR/export/mask_grounding_dino_model.onnx \\\n",
    "           results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate TensorRT engine using tao deploy\n",
    "!tao deploy mask_grounding_dino gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \\\n",
    "                               gen_trt_engine.onnx_file=$RESULTS_DIR/export/mask_grounding_dino_model.onnx \\\n",
    "                               gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/mask_grounding_dino_model.engine \\\n",
    "                               results_dir=$RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference with generated TensorRT engine\n",
    "!tao deploy mask_grounding_dino inference -e $SPECS_DIR/infer.yaml \\\n",
    "                              inference.trt_engine=$RESULTS_DIR/gen_trt_engine/mask_grounding_dino_model.engine \\\n",
    "                              results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"trt_inference\", \"images_annotated\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
